# -*- coding: utf-8 -*-
import pandas as pd
import numpy as np

import gc
import datetime

#导入模型相关函数
from sklearn.ensemble import RandomForestRegressor
from sklearn.ensemble import GradientBoostingRegressor

#导入模型保存、模型评估等函数
from sklearn.externals import joblib
from sklearn.metrics import r2_score
from sklearn.metrics import mean_squared_error   #均方误差回归损失

from multiprocessing import cpu_count

#划分测试和训练集
from sklearn.model_selection import train_test_split


from gen_train_data import get_train_data, get_desc_of_all_fea, data_standardization

#获取rf_regressor 模型
def train_model(train, feaList, labelName, model=None, saveToFile=None):
    #定义模型
    if model is None: #默认采用随机森林模型
        model = RandomForestRegressor(n_estimators=10, criterion='mse', max_depth=None, min_samples_split=2, 
                                  min_samples_leaf=1, min_weight_fraction_leaf=0.0, max_features='auto', max_leaf_nodes=None,
                                  min_impurity_decrease=0.0, min_impurity_split=None, 
                                  bootstrap=True, oob_score=False, n_jobs=4, random_state=None, verbose=0, warm_start=False)
    
    #训练集
    X_train = train[feaList].values
    y_train = train[[labelName]].values.ravel()
    
    # 拟合模型
    model.fit(X_train, y_train)
    
    print ('train model ok')
    
    # 保存模型
    if saveToFile is not None:
        joblib.dump(model, saveToFile)
        
        print ('model saveToFile!')
        
    #删除数据，并且释放内存
    del X_train, y_train
    gc.collect()
    
    return model

#测试对应模型，输出模型结果
def test_model(test, model, feaList, labelName):
    X_test = test[feaList].values
    y_test = test[[labelName]].values.ravel()
    
    y_test_pred = model.predict(X_test)
    
    # 输出测试结果
    mse = mean_squared_error(y_test, y_test_pred)
    r2Score = r2_score(y_test, y_test_pred)
    
    print('MSE test: %.3f' % (mse))
    print('R^2 test: %.3f' % (r2Score))
    
    #删除数据，并且释放内存
    del X_test, y_test, y_test_pred
    gc.collect()
    
    return mse, r2Score

if __name__ == '__main__':
	runTime = datetime.datetime.now().strftime('%Y%m%d%H%M')
	nrows = 10
	get_new_train = True
	get_new_model = True
	do_test_model = True

	srcTrainDataPath = '../data/tap_fun_train.csv'
	trainDataPath = '../train_data/train_data' + runTime + '.csv'
	doneTrainDataPath = '../train_data/train_data.csv'

	model_name = 'gbrt'

	#初始化用于训练模型的CPU核数
	n_jobs_for_model = 4
	if cpu_count() > 4:
		n_jobs_for_model = int(cpu_count() / 2)

	modelSavePath = '../model/' + model_name + runTime + '.pickle'
	doneModelPath = '../model/' + model_name + '.pickle'

	model_dict = {
	    'rf': RandomForestRegressor(n_estimators=100, criterion='mse', max_depth=None, min_samples_split=2, 
	                              min_samples_leaf=1, min_weight_fraction_leaf=0.0, max_features='auto', max_leaf_nodes=None,
	                              min_impurity_decrease=0.0, min_impurity_split=None, 
	                              bootstrap=True, oob_score=False, n_jobs=n_jobs_for_model, random_state=47, verbose=0, warm_start=False),

	    'gbrt':GradientBoostingRegressor(loss='ls', learning_rate=0.1, n_estimators=100, max_features='auto', max_depth=12, 
	                                     subsample=1.0,criterion='friedman_mse', min_samples_split=2, min_samples_leaf=1,
	                                     min_weight_fraction_leaf=0.0, min_impurity_decrease=0.0,
	                                     min_impurity_split=None, init=None, random_state=47, 
	                                     alpha=0.9, verbose=0, max_leaf_nodes=None, warm_start=False, presort='auto')
	}

	feaList = [
		#  'user_id',
		#  'register_time',
		'wood_add_value',
		'wood_reduce_value',
		'stone_add_value',
		'stone_reduce_value',
		'ivory_add_value',
		'ivory_reduce_value',
		'meat_add_value',
		'meat_reduce_value',
		'magic_add_value',
		'magic_reduce_value',
		'infantry_add_value',
		'infantry_reduce_value',
		'cavalry_add_value',
		'cavalry_reduce_value',
		'shaman_add_value',
		'shaman_reduce_value',
		'wound_infantry_add_value',
		'wound_infantry_reduce_value',
		'wound_cavalry_add_value',
		'wound_cavalry_reduce_value',
		'wound_shaman_add_value',
		'wound_shaman_reduce_value',
		'general_acceleration_add_value',
		'general_acceleration_reduce_value',
		'building_acceleration_add_value',
		'building_acceleration_reduce_value',
		'reaserch_acceleration_add_value',
		'reaserch_acceleration_reduce_value',
		'training_acceleration_add_value',
		'training_acceleration_reduce_value',
		'treatment_acceleraion_add_value',
		'treatment_acceleration_reduce_value',
		'bd_training_hut_level',
		'bd_healing_lodge_level',
		'bd_stronghold_level',
		'bd_outpost_portal_level',
		'bd_barrack_level',
		'bd_healing_spring_level',
		'bd_dolmen_level',
		'bd_guest_cavern_level',
		'bd_warehouse_level',
		'bd_watchtower_level',
		'bd_magic_coin_tree_level',
		'bd_hall_of_war_level',
		'bd_market_level',
		'bd_hero_gacha_level',
		'bd_hero_strengthen_level',
		'bd_hero_pve_level',
		'sr_scout_level',
		'sr_training_speed_level',
		'sr_infantry_tier_2_level',
		'sr_cavalry_tier_2_level',
		'sr_shaman_tier_2_level',
		'sr_infantry_atk_level',
		'sr_cavalry_atk_level',
		'sr_shaman_atk_level',
		'sr_infantry_tier_3_level',
		'sr_cavalry_tier_3_level',
		'sr_shaman_tier_3_level',
		'sr_troop_defense_level',
		'sr_infantry_def_level',
		'sr_cavalry_def_level',
		'sr_shaman_def_level',
		'sr_infantry_hp_level',
		'sr_cavalry_hp_level',
		'sr_shaman_hp_level',
		'sr_infantry_tier_4_level',
		'sr_cavalry_tier_4_level',
		'sr_shaman_tier_4_level',
		'sr_troop_attack_level',
		'sr_construction_speed_level',
		'sr_hide_storage_level',
		'sr_troop_consumption_level',
		'sr_rss_a_prod_levell',
		'sr_rss_b_prod_level',
		'sr_rss_c_prod_level',
		'sr_rss_d_prod_level',
		'sr_rss_a_gather_level',
		'sr_rss_b_gather_level',
		'sr_rss_c_gather_level',
		'sr_rss_d_gather_level',
		'sr_troop_load_level',
		'sr_rss_e_gather_level',
		'sr_rss_e_prod_level',
		'sr_outpost_durability_level',
		'sr_outpost_tier_2_level',
		'sr_healing_space_level',
		'sr_gathering_hunter_buff_level',
		'sr_healing_speed_level',
		'sr_outpost_tier_3_level',
		'sr_alliance_march_speed_level',
		'sr_pvp_march_speed_level',
		'sr_gathering_march_speed_level',
		'sr_outpost_tier_4_level',
		'sr_guest_troop_capacity_level',
		'sr_march_size_level',
		'sr_rss_help_bonus_level',
		'pvp_battle_count',
		'pvp_lanch_count',
		'pvp_win_count',
		'pve_battle_count',
		'pve_lanch_count',
		'pve_win_count',
		'avg_online_minutes',
		'pay_price',
		'pay_count',
		#  'prediction_pay_price',
		'year',
		'month',
		'day',
		'hour',
		'weekday',

		#各种材料使用情况
		'wood_usage',
		'stone_usage',
		'ivory_usage',
		'meat_usage',
		'magic_usage',
		'infantry_usage',
		'cavalry_usage',
		'shaman_usage',

		#各种材料使用情况分布
		'usage_max',
		'usage_min',
		'usage_mean',
		'usage_std',
		'usage_gap',
		'usage_cv',

		#恢复/消耗情况
		'wound_infantry_renew',
		'wound_cavalry_renew',
		'wound_shaman_renew',

		#恢复/消耗分布
		'renew_max',
		'renew_std',

		#加速使用情况
		'general_acceleration_acc_usage',
		'building_acceleration_acc_usage',
		'reaserch_acceleration_acc_usage',
		'training_acceleration_acc_usage',
		'treatment_acceleration_acc_usage',

		#加速使用情况分布
		'acc_usage_max',
		'acc_usage_min',
		'acc_usage_mean',
		'acc_usage_std',
		'acc_usage_gap',
		'acc_usage_cv',

		#建筑等级发展分布
		'bd_level_max',
		'bd_level_min',
		'bd_level_mean',
		'bd_level_std',
		'bd_level_gap',
		'bd_level_cv',

		#科研等级发展分布
		'sr_level_max',
		'sr_level_min',
		'sr_level_mean',
		'sr_level_std',
		'sr_level_gap',
		'sr_level_cv',

		#等级发展均衡性
		'max_level_balance',
		'mean_level_balance',

		#PVP/PVE主动发起比例
		'pvp_battle_lanch_ratio',
		'pve_battle_lanch_ratio',

		#PVP/PVE获胜比例
		'pvp_battle_win_ratio',
		'pve_battle_win_ratio',

		#主动发起的PVP/PVE获胜比例
		'pvp_battle_lanch_win_ratio',
		'pve_battle_lanch_win_ratio',

		#被动发起的PVP/PVE获胜比例
		'pvp_battle_no_lanch_win_ratio',
		'pve_battle_no_lanch_win_ratio'
	]

	labelName = 'prediction_pay_price'

	#初始化训练数据集
	df = None
	if get_new_train is True:
		df = get_train_data(srcTrainDataPath, nrows=nrows)

		#获取各个特征的统计值
		mean_dict, max_dict, min_dict, std_dict = get_desc_of_all_fea(df, feaList)

		#对数据集进行标准化处理
		df = data_standardization(df, feaList, mean_dict, max_dict, min_dict, std_dict)

		#将训练集存入文件
		df.to_csv(trainDataPath, index=None)
	else:
		df = pd.read_csv(doneTrainDataPath)

	print ('get_train_data ok!', datetime.datetime.now())

	#设置模型类型
	model = model_dict[model_name]
	if get_new_model is True: #需要重新训练模型，则重新训练模型
		#训练模型 
		df = df.fillna(0)
		oldtime=datetime.datetime.now()
		model = train_model(df, feaList, labelName, model=model, saveToFile=modelSavePath)
		newtime=datetime.datetime.now()
		print (model.get_params)
		print('train: old=', oldtime, ' new=', newtime, 'gap= %s'%(newtime - oldtime) )
	else:
		model = joblib.load(doneModelPath)

	#测试模型
	if do_test_model:
		oldtime=datetime.datetime.now()
		mse, r2Score = test_model(df, model, feaList, labelName)
		newtime=datetime.datetime.now()
		print('test: old=', oldtime, ' new=', newtime, 'gap= %s'%(newtime - oldtime) )


